#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rtabmap_msgs.srv import GetMap
import open3d as o3d
import numpy as np
# ros_numpy is a common and powerful library for converting between ROS messages and NumPy arrays.
# If this import fails, you may need to install it: sudo apt-get install ros-humble-ros-numpy
import ros_numpy

class MapSaver(Node):
    def __init__(self):
        super().__init__('map_saver_client')
        self.cli = self.create_client(GetMap, '/rtabmap/get_map')
        while not self.cli.wait_for_service(timeout_sec=1.0):
            self.get_logger().info('service not available, waiting again...')
        self.req = GetMap.Request()

    def send_request(self):
        # Request the global map, optimized, and ensure we get the point cloud data
        self.req.global_map = True
        self.req.optimized = True
        self.req.graph_only = False # This MUST be false to get the dense point cloud
        self.future = self.cli.call_async(self.req)
        rclpy.spin_until_future_complete(self, self.future)
        return self.future.result()

def main(args=None):
    rclpy.init(args=args)

    map_saver = MapSaver()
    response = map_saver.send_request()

    if response and response.data:
        map_saver.get_logger().info('Map data received, processing...')
        
        # The map data is in a ROS PointCloud2 format. We need to convert it.
        point_cloud2_msg = response.data.point_cloud
        
        if point_cloud2_msg.width == 0 or point_cloud2_msg.height == 0:
            map_saver.get_logger().error('Received map data, but it appears to be empty!')
            map_saver.destroy_node()
            rclpy.shutdown()
            return

        # Use ros_numpy to efficiently convert the PointCloud2 message to a NumPy array of XYZ points
        points_xyz = ros_numpy.point_cloud2.pointcloud2_to_xyz_array(point_cloud2_msg)

        # Create an Open3D point cloud object
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points_xyz)
        
        # Save as a PLY file, which is the best format for Three.js
        output_filename = "rtabmap_cloud.ply"
        o3d.io.write_point_cloud(output_filename, pcd)
        map_saver.get_logger().info(f'Successfully saved dense point cloud map to {output_filename}')
    else:
        map_saver.get_logger().error('Failed to call service /rtabmap/get_map or received no data')

    map_saver.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()